#!/usr/bin/env python

#############################################################################
# torrent.py: Contains the torrent class and other functions useful for
#             controlling and probing torrent behavior.
#############################################################################
# Copyright (C) 2012
#
# Authors: Dan Larkin <danielhlarkin@gmail.com>
#          Yonatan Naamad <cached@gmail.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#############################################################################

import logging
import bencode
import time
import socket
import os
import network
import binascii
import dcontrol
import shutil
import hashlib

_ERROR = -1
_DOWNLOADING = 0
_DOWNLOADED = 1
_HEX = 2
_BINARY = 3

def _compute_size(torrent_data):
    '''Gets the size of the PAYLOAD of the torrent file.'''
    metainfo = bencode.bdecode(torrent_data)
    return metainfo['info']['length']
    
def _compute_hash(torrent_data):
    '''Computes the hex hash of the torrent file.'''
    metainfo = bencode.bdecode(torrent_data)
    return hashlib.sha1(bencode.bencode(metainfo['info'])).digest()

def torrents_dir():
    from dcontrol import DControl
    return DControl.torrents_dir
    
def files_dir():
    from dcontrol import DControl
    return DControl.files_dir

def get_instance(t_hash, hash_type=_BINARY):
    '''Returns the instance of Torrent with the corresponding hash. Creates it
       if needed.'''
       
    if hash_type == _HEX:
        t_hash = binascii.unhexlify(t_hash)
        
    if t_hash in Torrent.t_map:
        return Torrent.t_map[t_hash]
    else:
        return Torrent(t_hash, hash_type=_BINARY)
    
def run_status_updater():
    '''Places .cpx_complete files in the directories containing
       completed downloads.'''
       
    def update_status():
        import time
        while 1:
            for i in dcontrol.get_torrents_info(keys=['is_finished']):
                debug.info("Finished status of torrent %s is %s" % (i.get_hex(), i['is_finished']))
                if i['is_finished']:
                    if not i.is_downloaded():
                        f = open(os.path.join(files_dir(), self.get_hex()) + "/.cpx_complete",'w')
                        f.write('True')
                        f.close()
            time.sleep(5)
    thread.start_new_thread(update_status,tuple())

def load_torrents():
    try:
        f = open(os.path.join(Torrent.data_dir, 'saved.txt'),'r')
        s = f.read().replace('\n','')[1:-1]
        f.close()
    except:
        return

    def success_callback(t):
        pass
    def failure_callback(t, err):
        logging.warning("Error downloading %s: '%s'. Removing from resume list." % (t, err))
        t.remove()

    for t in s.split('), ('):
        t_hex, min_alloc, max_alloc = t.split(', ')[0][1:-1], int(t.split(', ')[1]), int(t.split(', ')[2])

        if not binascii.unhexlify(t_hex) in Torrent.t_map:
            t = get_instance(t_hex, _HEX)
            t.acq_file(success_callback, failure_callback)
            t.add()
            t.set_allocations(min_alloc, max_alloc)

class Torrent(object):
    database_ip = None
    database_port = None
    dl_timeout = None
    data_dir = None
    client = None
    history_depth = None
    
    t_map = {}
    
    def __init__(self, t_hash, hash_type=_BINARY):
        if hash_type == _BINARY:
            self.hash = t_hash
        elif hash_type == _HEX:
            self.hash = binascii.unhexlify(t_hash)
            
        self.active = False
        self.min_alloc = 0
        self.max_alloc = 0
        self.next_request = 0
        
        self.payload_size = None
        self.recent_up_rates = [None for i in range(Torrent.history_depth)]
        
        Torrent.t_map[self.hash] = self
        
    def is_downloaded(self):
        '''Checks if the download of the torrent payload has completed.'''
        if torrents_dir() == None or files_dir() == None:
            raise LookupError("torrents_dir and files_dir need to be set first.")
        else:
            if '.cpx_complete' in os.listdir(os.path.join(files_dir(),self.get_hex())):
                return True
            else:
                return False
    
    def have_torrent_file(self):
        '''Checks if I already have the .torrent file.'''
        if torrents_dir() == None or files_dir() == None:
            raise LookupError("torrents_dir and files_dir need to be set first.")
        else:
            if self._make_name() in os.listdir(torrents_dir()):
                return True
            else:
                return False
    
    def torrent_file_downloading(self):
        '''Checks if I'm in the process of downloading the .torrent file.'''
        if torrents_dir() == None or files_dir() == None:
            raise LookupError("torrents_dir and files_dir need to be set first.")
        else:
            if self._make_temp_name() in os.listdir(torrents_dir()):
                return True
            else:
                return False
                
    def acq_file(self, success_callback = None, failure_callback = None):
        '''Downloads the .torrent file from the server, if necessary.
           You likely want to run this method in its own thread. Runs 
           success_callback(self) on success and failure_callback(self,msg)
           on failure. This is asynchronous, so nothing is returned immediately.'''
        
        def dl(success_callback, failure_callback):
            temp_path = os.path.join(torrents_dir(), self._make_temp_name())
            final_path = os.path.join(torrents_dir(), self._make_name())

            #check if I already have the file
            try:
                h = _compute_hash(open(final_path,'rb').read())
                if h == self.hash:
                    logging.info("No need to download a torrent I already have.")
                    success_callback(self)
                    return
            except Exception as e:
                pass

            t_name = self._make_temp_name()
            t_dir = torrents_dir()
        
            #check if I'm in the process of downloading the .torrent file
            if self.torrent_file_downloading():
                #the file exists. check that the download hasn't expired
            
                f = open(temp_path,"r")
                old_time = float(f.read())
                f.close()
            
                if time.time() - old_time > Torrent.dl_timeout:
                    #the download has expired. delete it.
                    logging.info("Old %s .torrent download is expired. Deleting." % self.get_hex())
                    os.unlink(temp_path)
                else:
                    return failure_callback(self, "Download already in progress.")

            #make the temporary file for locking
            f = open(temp_path,'w')
            f.write(str(time.time()))
            f.close()

            #connect to the torrent server
            logging.info("Connecting to torrent database to begin download.")
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.connect((Torrent.database_ip, Torrent.database_port))

            packet = network.write_schema({'T': self.get_bin()},network.GET_schema)
            logging.info("Sending packet requesting torrent %s" % self.get_hex())
            s.send(chr(network.GET_op))
            s.send(packet)

            #download the file
            op = s.recv(1)
            data = network.read_packet(s,network.GIVE_op,network.GIVE_schema)
            logging.info(".torrent file downloaded. Saving.")

            #write it to the .torrent file
            f = open(final_path,'wb')
            f.write(data['F'][0])
            f.close()

            #delete temp file
            os.unlink(temp_path)
        
            logging.info("Executing success callback for .torrent download.")
            success_callback(self)
        
        logging.info("Beginning blocking call to dl")
        dl(success_callback, failure_callback)
                             
    def add(self, up_throttle=None):
        '''Add a torrent to deluge.'''
        if self in Torrent.client.delete_queue:
            Torrent.client.delete_queue.remove(self)
            
        if Torrent.client.allocate_space(self):
            dcontrol.add_torrent(self, up_throttle=up_throttle, down_throttle = None)
            self.active = True
            Torrent.client.active_torrents.add(self)
        else:
            logging.error("Unable to add torrent due to space reservation issues.")
        
        self.jointime = time.time()
        self.save()
        
    def remove(self):
        '''Removes torrent from deluge and marks its disk space as deletable.'''
        dcontrol.remove_torrent(self)
        Torrent.client.active_torrents.remove(self)
        self.active = False
        Torrent.client.free_space(self)
        self.save()
        
    def delete_files(self):
        '''Deletes the torrent's entire download directory.'''
        shutil.rmtree(self, os.path.join(files_dir(),self.get_hex()))
        
    def set_allocations(self, min_alloc=None, max_alloc=None):
        '''Set the maximum and minimum allocations for the torrent, in KBps.'''
        if min_alloc != None:
            self.min_alloc = min_alloc
        if max_alloc != None:
            self.max_alloc = max_alloc
        self.set_up_throttle(max_alloc)
        self.save()
        
    def save(self):
        l = []
        for t in Torrent.t_map.values():
            if t.is_active():
                l.append((t.get_hex(), t.min_alloc, t.max_alloc))

        f = open(os.path.join(Torrent.data_dir, 'saved.txt'),'w')
        f.write(str(l)[1:-1])
        f.close()

    def __str__(self):
        return "Torrent<%s, min=%d, max=%d>" % (self.get_hex(), self.min_alloc, self.max_alloc)

    def __repr__(self):
        return self.__str__()

    def set_up_throttle(self, throttle):
        '''Sets the upload throttle of the torrent.'''
        dcontrol.set_up_throttle(self, throttle)
        
    def set_down_throttle(self, throttle):
        '''Sets the download throttle of the torrent.'''
        dcontrol.set_down_throttle(self, throttle)
        
    def is_active(self):
        '''Return whether the torrent is active or has been removed.'''
        return self.active
        
    def get_file(self):
        '''Returns the path the the .torrent file.'''
        return os.path.join(torrents_dir(), self.get_hex() + ".torrent")
        
    def get_upload_rate(self):
        '''Returns my current upload rate, in KBps.'''
        return dcontrol.get_upload_rate(self)
        
    def get_payload_size(self):
        '''Returns the size of my payload in bytes.'''
        if self.payload_size == None:
            s = open(self.get_file(),'r').read()
            self.payload_size = _compute_size(s)
            
        return self.payload_size
        
    def get_hex(self):
        '''Returns the hex value of the hash.'''
        return binascii.hexlify(self.hash)
        
    def get_bin(self):
        '''Returns the binary value of the hash.'''
        return self.hash

    def _make_temp_name(self):
        return self.get_hex() + ".torrent.part"
                
    def _make_name(self):
        return self.get_hex() + ".torrent"
        
def _timeout(func, args=(), kwargs={}, timeout_duration=1):
    '''Runs a function, but times it out after timeout_duration seconds.'''
    import threading
    class InterruptableThread(threading.Thread):
        def __init__(self):
            threading.Thread.__init__(self)
            self.result = None

        def run(self):
            try:
                self.result = func(*args, **kwargs)
            except:
                self.result = None
    it = InterruptableThread()
    it.start()
    it.join(timeout_duration)
        
def _nonblocking_timeout(func, args=(), kwargs={}, timeout_duration=1):
    import thread
    thread.start_new_thread(_timeout,(func, args, kwargs, timeout_duration))
